{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/5\n",
      "20/20 [==============================] - 7s 186ms/step - factorized_top_k/top_1_categorical_accuracy: 1.0000e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0015 - factorized_top_k/top_10_categorical_accuracy: 0.0045 - factorized_top_k/top_50_categorical_accuracy: 0.0395 - factorized_top_k/top_100_categorical_accuracy: 0.0899 - loss: 33381.9342 - regularization_loss: 0.0000e+00 - total_loss: 33381.9342\n",
      "Epoch 2/5\n",
      "20/20 [==============================] - 4s 183ms/step - factorized_top_k/top_1_categorical_accuracy: 1.6250e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0051 - factorized_top_k/top_10_categorical_accuracy: 0.0152 - factorized_top_k/top_50_categorical_accuracy: 0.1069 - factorized_top_k/top_100_categorical_accuracy: 0.2090 - loss: 31094.9421 - regularization_loss: 0.0000e+00 - total_loss: 31094.9421\n",
      "Epoch 3/5\n",
      "20/20 [==============================] - 5s 188ms/step - factorized_top_k/top_1_categorical_accuracy: 4.2500e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0099 - factorized_top_k/top_10_categorical_accuracy: 0.0263 - factorized_top_k/top_50_categorical_accuracy: 0.1595 - factorized_top_k/top_100_categorical_accuracy: 0.2870 - loss: 30341.3872 - regularization_loss: 0.0000e+00 - total_loss: 30341.3872\n",
      "Epoch 4/5\n",
      "20/20 [==============================] - 4s 185ms/step - factorized_top_k/top_1_categorical_accuracy: 5.6250e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0124 - factorized_top_k/top_10_categorical_accuracy: 0.0345 - factorized_top_k/top_50_categorical_accuracy: 0.2025 - factorized_top_k/top_100_categorical_accuracy: 0.3465 - loss: 29903.2236 - regularization_loss: 0.0000e+00 - total_loss: 29903.2236\n",
      "Epoch 5/5\n",
      "20/20 [==============================] - 4s 185ms/step - factorized_top_k/top_1_categorical_accuracy: 4.1250e-04 - factorized_top_k/top_5_categorical_accuracy: 0.0131 - factorized_top_k/top_10_categorical_accuracy: 0.0379 - factorized_top_k/top_50_categorical_accuracy: 0.2287 - factorized_top_k/top_100_categorical_accuracy: 0.3839 - loss: 29586.2553 - regularization_loss: 0.0000e+00 - total_loss: 29586.2553\n",
      "5/5 [==============================] - 2s 172ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0000e+00 - factorized_top_k/top_5_categorical_accuracy: 1.5000e-04 - factorized_top_k/top_10_categorical_accuracy: 0.0018 - factorized_top_k/top_50_categorical_accuracy: 0.0578 - factorized_top_k/top_100_categorical_accuracy: 0.1513 - loss: 33161.6159 - regularization_loss: 0.0000e+00 - total_loss: 33161.6159\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'factorized_top_k/top_1_categorical_accuracy': 0.0,\n",
       " 'factorized_top_k/top_5_categorical_accuracy': 0.0001500000071246177,\n",
       " 'factorized_top_k/top_10_categorical_accuracy': 0.0017999999690800905,\n",
       " 'factorized_top_k/top_50_categorical_accuracy': 0.057750001549720764,\n",
       " 'factorized_top_k/top_100_categorical_accuracy': 0.15129999816417694,\n",
       " 'loss': 30168.92578125,\n",
       " 'regularization_loss': 0,\n",
       " 'total_loss': 30168.92578125}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Dict, Text\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow_recommenders as tfrs\n",
    "\n",
    "# 加载电影评分数据集\n",
    "ratings = tfds.load('movielens/100k-ratings', split=\"train\")\n",
    "# 加载电影数据集\n",
    "movies = tfds.load('movielens/100k-movies', split=\"train\")\n",
    "\n",
    "# 将电影评分数据中的电影ID和用户ID转换为数字\n",
    "ratings = ratings.map(lambda x: {\n",
    "    \"movie_id\": tf.strings.to_number(x[\"movie_id\"]),\n",
    "    \"user_id\": tf.strings.to_number(x[\"user_id\"])\n",
    "})\n",
    "# 将电影数据中的电影ID转换为数字\n",
    "movies = movies.map(lambda x: tf.strings.to_number(x[\"movie_id\"]))\n",
    "\n",
    "class Model(tfrs.Model):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # 定义用户嵌入层\n",
    "        self.user_model = tf.keras.layers.Embedding(\n",
    "            input_dim=2000, output_dim=64)\n",
    "        # 定义电影嵌入层\n",
    "        self.item_model = tf.keras.layers.Embedding(\n",
    "            input_dim=2000, output_dim=64)\n",
    "        # 定义检索任务，使用FactorizedTopK指标\n",
    "        self.task = tfrs.tasks.Retrieval(\n",
    "            metrics=tfrs.metrics.FactorizedTopK(\n",
    "                candidates=movies.batch(128).map(self.item_model)\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -> tf.Tensor:\n",
    "        # 计算用户嵌入向量\n",
    "        user_embeddings = self.user_model(features[\"user_id\"])\n",
    "        # 计算电影嵌入向量\n",
    "        movie_embeddings = self.item_model(features[\"movie_id\"])\n",
    "        # 计算损失\n",
    "        return self.task(user_embeddings, movie_embeddings)\n",
    "\n",
    "# 创建模型实例\n",
    "model = Model()\n",
    "# 编译模型，指定优化器\n",
    "model.compile(optimizer=tf.keras.optimizers.Adagrad(0.5))\n",
    "# 设置随机种子\n",
    "tf.random.set_seed(42)\n",
    "# 对评分数据进行洗牌\n",
    "shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)\n",
    "# 将洗牌后的数据集分为训练集和测试集\n",
    "train = shuffled.take(80_000)\n",
    "test = shuffled.skip(80_000).take(20_000)\n",
    "\n",
    "# 训练模型\n",
    "model.fit(train.batch(4096), epochs=5)\n",
    "# 评估模型性能\n",
    "model.evaluate(test.batch(4096), return_dict=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
